/*!
  \file
  Header file for an allocator which uses cudaHostAlloc
*/

#ifndef SCIGPU_HHV_CUDA_HOST_ALLOCATOR_HPP
#define SCIGPU_HHV_CUDA_HOST_ALLOCATOR_HPP

#include <new>
#include <cuda_runtime.h>

// ==============================


/*
  The core of this code was copied from
  http://www.drdobbs.com/184403759
*/

namespace SciGPU {
  namespace HHV {

    //! Templated allocator for pinned host memory
    /*!
      This is an object type suitable for use as an \c allocator
      by STL containers.
      Instead of \c new and \c delete, memory is managed using
      \c cudaHostAlloc and \c cudaFreeHost, so that pinned
      memory can be used to buffer CUDA objects.
      The core of this code was copied from
      <a href="http://www.drdobbs.com/184403759">Dr Dobbs</a>.
      The flags used by \c cudaHostAlloc are set via the
      template parameter \a flags.
      @tparam T     The data type to allocate
      @tparam flags The flags to be passed to \c cudaHostAlloc
    */
    template <typename T, unsigned int flags=cudaHostAllocDefault> class cuda_host_allocator {
    public:
      //! Export of the data type, \a T, allocated
      typedef T                 value_type;
      //! Export of the pointer type
      typedef value_type*       pointer;
      //! Export of the \c const pointer type
      typedef const value_type* const_pointer;
      //! Export of the reference type
      typedef value_type&       reference;
      //! Export of the \c const reference type
      typedef const value_type& const_reference;
      //! Export of the size type
      typedef std::size_t       size_type;
      //! Export of the difference type
      typedef std::ptrdiff_t    difference_type;
  
      template <typename U> 
      struct rebind {
        typedef cuda_host_allocator<U,flags> other;
      };

      //! Default constructor
      cuda_host_allocator( void ) {}
      //! Copy Constructor
      cuda_host_allocator( const cuda_host_allocator& src ) {}
      //! Constructor from different allocator
      template <typename U> 
      cuda_host_allocator( const cuda_host_allocator<U,flags>& src ) {}

      //! Destructor
      ~cuda_host_allocator() {}

      //! Method to return address
      pointer address( reference x ) const {
        return &x;
      }
      //! Method to return a constant address
      const_pointer address( const_reference x ) const { 
        return x;
      }

      //! Allocation method
      pointer allocate( size_type n, const_pointer = 0 ) {
        void* p;
        cudaError_t err = cudaHostAlloc( &p, n * sizeof(T), flags );
        if( cudaSuccess != err ) {
          throw std::bad_alloc();
        }
        return( static_cast<pointer>(p) );
      }

      //! Deallocation method
      void deallocate( pointer p, size_type s ) {
        s = 10; // Silence compiler
        cudaError_t err = cudaFreeHost( p );
        if( cudaSuccess != err ) {
          throw std::bad_alloc();
        }
      }

      size_type max_size( void ) const { 
        return( static_cast<size_type>(-1) / sizeof(T) );
      }

      //! Create a new object at location
      void construct( pointer p, const value_type& x ) { 
        new(p) value_type(x); 
      }

      //! Destroy an object
      void destroy( pointer p ) {
        p->~value_type();
      }

    private:
      //! Copy operator (inhibited)
      void operator=(const cuda_host_allocator& src);
    };


    //! Specialisation for voids
    template<> class cuda_host_allocator<void>
    {
      typedef void        value_type;
      typedef void*       pointer;
      typedef const void* const_pointer;

      template <class U> 
      struct rebind { typedef cuda_host_allocator<U> other; };
    };
  }
}


//! Equality comparison
template <typename T, unsigned int f1, unsigned int f2>
inline bool operator==( const SciGPU::HHV::cuda_host_allocator<T,f1>& a, 
                        const SciGPU::HHV::cuda_host_allocator<T,f2>& b ) {
  return( f1 == f2 );
}

//! Inequality comparison
template <typename T, unsigned int f1, unsigned int f2>
inline bool operator!=(const SciGPU::HHV::cuda_host_allocator<T,f1>& a, 
                       const SciGPU::HHV::cuda_host_allocator<T,f2>& b ) {
  return( f1 != f2 );
}

#endif
